import torch.nn as nn
from .encoder import Encoder
from .decoder import Decoder

filter = [64, 64, 128, 256, 512]


class Net(nn.Module):
    def __init__(self, input_channels=1, filters=None, output_dementions=1, bias=True):
        super().__init__()
        if filters is None:
            filters = filter
        self.input_channels = input_channels
        self.output_dementions = output_dementions
        self.encoder = Encoder(filters, input_channels, bias=bias)
        self.decoder = Decoder(filters, output_dementions, bias=bias)
        # self.drop_path_prob = 0

    def forward(self, input):
        outEncode = self.encoder(input)
        outDecode = self.decoder(outEncode)

        return outDecode


if __name__ == '__main__':
    import torch

    x = torch.autograd.Variable(torch.randn(32, 1, 224, 184))
    model = Net()
    yy = model.forward(x)
    print(yy.shape)
